> ## Documentation Index
> Fetch the complete documentation index at: https://mintlify.com/mlfoundations/open_clip/llms.txt
> Use this file to discover all available pages before exploring further.

# Distributed Training

> Advanced techniques for efficient large-scale CLIP training

## Overview

OpenCLIP provides advanced distributed training techniques that enable efficient training on hundreds or thousands of GPUs. This guide covers memory-efficient distributed loss computation, gradient accumulation, mixed precision training, and performance optimizations.

## Memory-Efficient Distributed Loss

The standard CLIP contrastive loss requires computing a logit matrix of size `(batch_size × num_gpus) × (batch_size × num_gpus)`, leading to O(n²) memory complexity. For large-scale training, this becomes a bottleneck.

### The Problem: O(n²) Memory Complexity

**Without optimization:**

```
Global batch size: 256 samples/GPU × 128 GPUs = 32,768 samples
Logit matrix size: 32,768 × 32,768 = 1,073,741,824 elements
Memory (fp32): ~4 GB just for the logit matrix
```

This scales poorly and limits the number of GPUs you can use.

### The Solution: Local Loss with Gradient Gathering

OpenCLIP implements an efficient distributed loss computation that achieves O(n) memory complexity while maintaining identical numerical results.

```bash theme={null}
python -m open_clip_train.main \
    --local-loss \
    --gather-with-grad \
    # ... other arguments
```

**With optimization:**

```
Memory per GPU: O(local_batch_size) instead of O(global_batch_size²)
Numerical results: Identical to the naïve all-gather approach
```

### How It Works

1. **`--local-loss`**: Compute loss locally with gathered features, avoiding full global logit matrix
2. **`--gather-with-grad`**: Enable gradient flow through the feature gathering operation

Together, these flags enable linear memory scaling:

<CodeGroup>
  ```bash Single-Node (4 GPUs) theme={null}
  torchrun --nproc_per_node 4 -m open_clip_train.main \
      --train-data '/data/train.tar' \
      --batch-size 256 \
      --local-loss \
      --gather-with-grad \
      --model ViT-B-32
  ```

  ```bash Multi-Node (128 GPUs) theme={null}
  srun python -u src/open_clip_train/main.py \
      --train-data '/data/train.tar' \
      --batch-size 256 \
      --local-loss \
      --gather-with-grad \
      --model ViT-L-14
  # Effective batch: 256 × 128 = 32,768
  # Memory: O(256) per GPU, not O(32,768²)
  ```
</CodeGroup>

<Warning>
  Always use `--local-loss` and `--gather-with-grad` together for multi-node training (8+ GPUs). These flags are essential for scaling beyond small clusters.
</Warning>

## Gradient Accumulation

Gradient accumulation simulates larger batch sizes by accumulating gradients over multiple forward passes before updating weights.

### Basic Usage

```bash theme={null}
python -m open_clip_train.main \
    --batch-size 128 \
    --accum-freq 4 \
    # ... other arguments
```

**Effective batch size:**

```
effective_batch = batch_size × num_gpus × accum_freq
                = 128 × 8 × 4
                = 4,096
```

### When to Use Gradient Accumulation

Use `--accum-freq` when:

1. ✅ You need larger effective batch sizes than GPU memory allows
2. ✅ You want to maintain a specific batch size across different hardware
3. ✅ You're experimenting with very large batch sizes (>100k)

**However, try these first:**

1. Enable `--precision amp` (mixed precision)
2. Use `--grad-checkpointing` (memory-compute tradeoff)
3. Use `--local-loss --gather-with-grad` (for distributed training)

### Performance Implications

<Info>
  **Speed:** Samples/sec remains approximately constant

  * Without accumulation: Process 1024 samples in 1 step
  * With `--accum-freq 4`: Process 256 samples × 4 in 4 steps
  * Net throughput: Similar

  **Memory:** Additional GPU memory required for:

  * Cached features from all accumulated batches
  * Multiple loss computations (one per accumulated batch)
</Info>

### Real-World Example

```bash theme={null}
# Training ViT-L-14 on 8 A100 GPUs (40GB)
# Without gradient accumulation: OOM at batch_size 256
# With gradient accumulation: Works!

torchrun --nproc_per_node 8 -m open_clip_train.main \
    --train-data '/data/laion400m/{00000..41455}.tar' \
    --batch-size 128 \
    --accum-freq 2 \
    --precision amp \
    --grad-checkpointing \
    --model ViT-L-14 \
    --local-loss \
    --gather-with-grad
# Effective batch: 128 × 8 × 2 = 2,048
```

### Implementation Details

Gradient accumulation in OpenCLIP:

1. **First N-1 steps:** Forward pass with `torch.no_grad()`, cache features
2. **Nth step:** Re-run forward passes with gradients enabled
3. **Backward:** Compute gradients using cached features as negatives
4. **Step:** Update optimizer

References:

* [Cui et al. 2021](https://arxiv.org/abs/2112.09331)
* [Pham et al. 2021](https://arxiv.org/abs/2111.10050)

## Mixed Precision Training

Mixed precision training uses lower precision (float16 or bfloat16) for most computations while maintaining float32 for critical operations.

### Automatic Mixed Precision (AMP)

```bash theme={null}
# FP16 with automatic loss scaling (recommended)
python -m open_clip_train.main \
    --precision amp \
    # ... other arguments

# BFloat16 (if supported by hardware)
python -m open_clip_train.main \
    --precision amp_bf16 \
    # ... other arguments
```

**Benefits:**

* 🚀 **Speed:** 2-3× faster training on modern GPUs (A100, H100)
* 💾 **Memory:** \~50% reduction in activation memory
* 📊 **Accuracy:** Negligible impact with automatic loss scaling

### Precision Options

| Precision  | Description                      | Use Case                              |
| ---------- | -------------------------------- | ------------------------------------- |
| `amp`      | Automatic Mixed Precision (FP16) | **Recommended** for most training     |
| `amp_bf16` | AMP with BFloat16                | A100/H100 GPUs, more stable than FP16 |
| `fp32`     | Full 32-bit precision            | Debugging, baseline comparison        |
| `fp16`     | Pure FP16 (not recommended)      | Legacy, requires manual tuning        |
| `bf16`     | Pure BFloat16                    | Experimental                          |

### Hardware Support

**NVIDIA GPUs:**

* **Volta (V100):** FP16 via `--precision amp`
* **Ampere (A100):** FP16 or BF16 via `--precision amp` or `--precision amp_bf16`
* **Hopper (H100):** BF16 recommended via `--precision amp_bf16`

**AMD GPUs:**

* **MI250X:** FP16 via `--precision amp`

### Example: Mixed Precision Training

```bash theme={null}
torchrun --nproc_per_node 4 -m open_clip_train.main \
    --train-data '/data/train.tar' \
    --batch-size 512 \
    --precision amp \
    --model ViT-B-32
# Batch size doubled thanks to reduced memory usage
```

<Warning>
  Avoid using `--precision fp16` (pure FP16) without automatic loss scaling. Use `--precision amp` instead, which handles loss scaling automatically.
</Warning>

## Patch Dropout for Vision Transformers

Patch dropout randomly drops image patches during training, leading to 2-3× speedup for Vision Transformer models without accuracy loss.

### Research Background

[Li et al. 2022](https://arxiv.org/abs/2212.00794) showed that dropping 50-75% of visual tokens during training:

* ✅ Speeds up training by 2-3×
* ✅ Maintains final accuracy
* ✅ Acts as a form of data augmentation

### Usage

Set patch dropout in your model config or via command-line:

```bash theme={null}
# Enable 50% patch dropout during training
python -m open_clip_train.main \
    --model ViT-B-32 \
    --force-patch-dropout 0.5 \
    --epochs 30 \
    # ... other arguments
```

### Fine-tuning Without Patch Dropout

The paper recommends fine-tuning without patch dropout at the end:

```bash theme={null}
# Phase 1: Train with patch dropout (fast)
python -m open_clip_train.main \
    --model ViT-B-32 \
    --force-patch-dropout 0.5 \
    --epochs 28 \
    --name "phase1-with-dropout"

# Phase 2: Fine-tune without patch dropout (accuracy)
python -m open_clip_train.main \
    --model ViT-B-32 \
    --force-patch-dropout 0.0 \
    --resume /path/to/phase1/epoch_28.pt \
    --epochs 32 \
    --name "phase2-no-dropout"
```

### Recommended Values

| Model Size | Patch Dropout | Training Speedup |
| ---------- | ------------- | ---------------- |
| ViT-B/32   | 0.5           | \~2×             |
| ViT-B/16   | 0.5           | \~2×             |
| ViT-L/14   | 0.5-0.75      | \~2-3×           |
| ViT-H/14   | 0.75          | \~3×             |

<Info>
  Patch dropout only applies to Vision Transformer models. It has no effect on ResNet or ConvNext models.
</Info>

## Gradient Checkpointing

Gradient checkpointing trades compute for memory by recomputing activations during the backward pass instead of storing them.

```bash theme={null}
python -m open_clip_train.main \
    --grad-checkpointing \
    # ... other arguments
```

**Tradeoffs:**

* ✅ **Memory:** 30-50% reduction in activation memory
* ❌ **Speed:** \~20% slower due to recomputation
* ✅ **Batch Size:** Allows larger batch sizes

**When to use:**

* Training large models (ViT-L, ViT-H, ViT-bigG)
* Out of memory errors even with mixed precision
* Prefer larger batch sizes over speed

**Example:**

```bash theme={null}
# ViT-H-14 requires gradient checkpointing on A100-40GB
torchrun --nproc_per_node 8 -m open_clip_train.main \
    --model ViT-H-14 \
    --batch-size 64 \
    --precision amp \
    --grad-checkpointing \
    --local-loss \
    --gather-with-grad
```

## SyncBatchNorm

Synchronize batch normalization statistics across GPUs for models with BatchNorm layers.

```bash theme={null}
python -m open_clip_train.main \
    --use-bn-sync \
    # ... other arguments (for models with BatchNorm)
```

**When to use:**

* ResNet models with BatchNorm layers
* Small batch sizes per GPU (\<32)

**Not needed for:**

* Vision Transformers (use LayerNorm)
* Large batch sizes (>128 per GPU)

## Combining Techniques

### Small-Scale Training (1-4 GPUs)

```bash theme={null}
torchrun --nproc_per_node 4 -m open_clip_train.main \
    --train-data '/data/train.tar' \
    --batch-size 256 \
    --precision amp \
    --model ViT-B-32
# Simple and effective
```

### Medium-Scale Training (8-32 GPUs)

```bash theme={null}
torchrun --nproc_per_node 8 -m open_clip_train.main \
    --train-data '/data/train.tar' \
    --batch-size 128 \
    --precision amp \
    --local-loss \
    --gather-with-grad \
    --model ViT-L-14
# Add distributed loss optimization
```

### Large-Scale Training (64+ GPUs)

```bash theme={null}
srun python -u src/open_clip_train/main.py \
    --train-data '/data/train.tar' \
    --batch-size 256 \
    --precision amp_bf16 \
    --grad-checkpointing \
    --force-patch-dropout 0.5 \
    --local-loss \
    --gather-with-grad \
    --model ViT-H-14
# All optimizations enabled
```

### Very Large-Scale (256-1024 GPUs)

```bash theme={null}
srun python -u src/open_clip_train/main.py \
    --train-data '/data/laion-2b/{00000..100000}.tar' \
    --train-num-samples 2000000000 \
    --dataset-type webdataset \
    --dataset-resampled \
    --batch-size 256 \
    --precision amp_bf16 \
    --grad-checkpointing \
    --force-patch-dropout 0.75 \
    --local-loss \
    --gather-with-grad \
    --workers 12 \
    --model ViT-bigG-14 \
    --remote-sync s3://bucket/checkpoints \
    --delete-previous-checkpoint
# Maximum optimization for largest scale
```

## DDP Static Graph

PyTorch 1.11+ supports static graph optimization for DistributedDataParallel:

```bash theme={null}
python -m open_clip_train.main \
    --ddp-static-graph \
    # ... other arguments
```

**Benefits:**

* Slightly faster gradient synchronization
* Lower memory overhead

**Requirements:**

* Model architecture doesn't change during training
* PyTorch >= 1.11

## Torch Compile

PyTorch 2.0+ supports model compilation for faster execution:

```bash theme={null}
python -m open_clip_train.main \
    --torchcompile \
    # ... other arguments
```

**Benefits:**

* 10-30% speedup on A100/H100 GPUs
* Automatic kernel fusion and optimization

**Considerations:**

* First epoch is slower (compilation time)
* Requires PyTorch >= 2.0
* May have compatibility issues with some models

<Warning>
  If using `--grad-checkpointing` with `--torchcompile` and DDP, the DDP dynamo optimizer is automatically disabled to avoid compatibility issues.
</Warning>

## Int8 Training (Experimental)

OpenCLIP has beta support for int8 training using bitsandbytes:

```bash theme={null}
pip install bitsandbytes triton

python -m open_clip_train.main \
    --use-bnb-linear SwitchBackLinearGlobal \
    # ... other arguments
```

**Benefits:**

* \~10% training speedup for ViT-Huge
* Reduced memory usage
* No accuracy loss (preliminary results)

**Status:** Experimental, see [tutorial](https://github.com/mlfoundations/open_clip/blob/main/tutorials/int8_tutorial.ipynb)

## Performance Monitoring

### GPU Utilization

Monitor GPU usage during training:

```bash theme={null}
watch -n 1 nvidia-smi
```

**Target:** 90-100% GPU utilization

**If GPU utilization is low (\<80%):**

1. Increase `--workers` (data loading parallelism)
2. Use faster storage (NVMe SSD)
3. Increase `--batch-size` if memory allows
4. Profile data loading pipeline

### Throughput Measurement

OpenCLIP logs samples/sec during training:

```
Train Epoch: 0 [1280/10968539 (0%)] Batch (t): 0.850, 1506/s, 376/s/gpu
```

* **1506/s:** Global samples per second (all GPUs)
* **376/s/gpu:** Samples per second per GPU

**Typical values on A100 (40GB):**

| Model    | Batch/GPU | Precision | Samples/sec/GPU |
| -------- | --------- | --------- | --------------- |
| RN50     | 512       | amp       | \~1000          |
| ViT-B/32 | 320       | amp       | \~600           |
| ViT-L/14 | 128       | amp       | \~150           |
| ViT-H/14 | 64        | amp       | \~50            |

### Communication Overhead

For multi-node training, monitor network bandwidth:

```bash theme={null}
iftop -i ib0  # InfiniBand
iftop -i eth0 # Ethernet
```

**Expect:** High bandwidth during gradient synchronization (every step)

## Best Practices Summary

### ✅ Always Use

1. **Mixed Precision:** `--precision amp` (or `amp_bf16` on A100/H100)
2. **Distributed Loss:** `--local-loss --gather-with-grad` (for 8+ GPUs)
3. **WebDataset:** For datasets >10M samples

### ✅ Use When Needed

4. **Gradient Checkpointing:** `--grad-checkpointing` (for large models)
5. **Patch Dropout:** `--force-patch-dropout 0.5` (for ViT models)
6. **Gradient Accumulation:** `--accum-freq N` (when other options exhausted)

### ✅ Optimize For Your Setup

7. **Workers:** `--workers 4-12` (match to CPU cores per GPU)
8. **Batch Size:** Maximize per GPU (limited by memory)
9. **Save Frequency:** `--save-frequency 1` (or less frequent for large models)

### ❌ Avoid

* Pure FP16: Use `--precision amp` instead
* Small batches: \<64 per GPU reduces efficiency
* Too many workers: >16 per GPU causes overhead
* CSV format for large datasets: Use WebDataset

## Next Steps

<CardGroup cols={2}>
  <Card title="Multi-Node Training" icon="network-wired" href="/training/multi-node">
    Scale training across multiple machines with SLURM
  </Card>

  <Card title="Configuration" icon="sliders" href="/training/configuration">
    Explore all training configuration options
  </Card>

  <Card title="Data Preparation" icon="database" href="/training/data-preparation">
    Prepare datasets for efficient distributed training
  </Card>

  <Card title="Fine-tuning" icon="arrows-rotate" href="/training/fine-tuning">
    Fine-tune pretrained models with distributed techniques
  </Card>
</CardGroup>